"""
基于predict_iter.
在此基础上，用GA算法替代之前的迭代过程。其余不变
"""
import random
import pandas as pd
import csv
import numpy as np
from utils.base_util import new_fold, get_list_from_dict

from utils.painter import draw_hist, draw_line, draw_distribution_histogram
from network.ResNet_model import resnet50
from network.Basic_model import basic_net
from network.Fuse_model_1 import F1Net
from network.metrics import vm_sample
import torch
import utils.image_process as imp
from DRR.drr_generator import Projector
from configs.predict_conf import load_predict_config
from configs.config import global_var
from utils.label_transform import label2real
import matplotlib.pyplot as plt
from tqdm import tqdm
from utils.data_statistics import ResAnalyse
from utils.label_transform import rt2noised_rt
from numpy import mean, var

# 导入xy坐标
coord = np.load("data/coord/normed_coord_xy.npy")


class Predict:
    def __init__(self, mode, num_classes, model_path, CT_path, CT_pt_path, saving_drr_path):
        self.mode = mode
        # self.model = resnet50(num_classes=num_classes)
        # self.model = F1Net(in_channels=3)
        self.model = basic_net(in_channels=3, out_channels=num_classes)
        self.model.load_state_dict(torch.load(model_path))
        # if torch.cuda.is_available():
        #     self.model.cuda()
        # 导入用于DRR生成的对象
        self.data_generator = Projector(directory=CT_path)
        # 导入CT点云模型
        self.CT_pt = np.load(CT_pt_path)
        # 图片保存位置
        self.saving_drr_path = saving_drr_path

    def get_drr(self, pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz):
        new_drr = imp.get_new_drr_numpy(pre_alpha, pre_beta, pre_theta,
                                        pre_tx, pre_ty, pre_tz, "标准正位", self.data_generator)
        return new_drr[np.newaxis, :]

    def predict(self, test_img):
        self.model.eval()
        with torch.no_grad():
            # test_img = test_img.reshape(1, 1, 128, 128) / 255
            # img_tensor = torch.from_numpy(test_img).float()
            out = self.model(test_img)
        pre_r = out[:, 0: 3].detach().cpu().numpy()
        pre_t = out[:, 3: 6].detach().cpu().numpy()
        # 获取预测的DRR
        pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz = label2real(pre_r, pre_t, 1)
        return np.array([float(pre_alpha), float(pre_beta), float(pre_theta)]), \
            np.array([float(pre_tx), float(pre_ty), float(pre_tz)])


def add_coord(drr):
    # 对图像本身做归一化
    img_numpy = imp.norm_image(drr.astype(np.float64)) / 255
    # 添加xy坐标
    img_numpy = np.vstack((img_numpy, coord))
    img_tensor = torch.from_numpy(img_numpy).float()
    return img_tensor.unsqueeze(0)


class Pre:
    def __init__(self, _id):
        self.id = _id
        self.pre_data = pd.DataFrame(columns=['pre_rx', 'pre_ry', 'pre_rz', 'pre_tx', 'pre_ty', 'pre_tz'])
        self.error = pd.DataFrame(columns=['err_rx', 'err_ry', 'err_rz', 'err_tx', 'err_ty', 'err_tz'])
        self.VM = pd.DataFrame(columns=['VM'])


def para_limit(drr_para, _min=-50, _max=50):
    drr_para[3:] = np.clip(drr_para[3:], _min, _max)
    return drr_para


def best_visualization():
    """
    仅在本文件中有用，优化结构用
    :return:
    """
    # 可视化
    this_tru = np.array([rx, ry, rz, tx, ty, tz]).ravel()
    this_pre_1 = min_pre_drr_iter
    plt.figure(figsize=(8, 5))
    plt.subplot(1, 2, 1)
    plt.title('实际值'
              + "\n" +
              '角度' + str(np.around(this_tru[0:3], 2))
              + "\n" +
              '平移' + str(np.around(this_tru[3:], 2)))
    plt.imshow(test_drr.reshape(128, 128), cmap='gray')
    plt.subplot(1, 2, 2)
    plt.title('预测值'
              + "\n" +
              '角度' + str(np.around(this_pre_1[0:3], 2))
              + "\n" +
              '平移' + str(np.around(this_pre_1[3:], 2)))
    plt.imshow(best_pre_drr.reshape(128, 128), cmap='gray')
    plt.suptitle("{0}:验证{1}加噪的精度".format(epoch, img_name))
    if min_VM > 1:
        plt.savefig(predictor.saving_drr_path + "result_1+/第{0}组_VM_{1}.png".format(epoch, round(min_VM, 3)))
    else:
        plt.savefig(predictor.saving_drr_path + "result_1-/第{0}组_VM_{1}.png".format(epoch, round(min_VM, 3)))
    # plt.show()
    plt.close()


if __name__ == "__main__":
    new_fold('pre_result_GA', ['result_1+', 'result_1-', 'statistical_result', 'VM_comparison',
                               'datasheet/pre_iter', 'datasheet/error_iter', 'datasheet/VM_iter'])
    # 导入相应的模型
    pred_conf = load_predict_config(global_var.loading_address)
    predictor = Predict(**pred_conf)
    # 添加测试轮数
    epochs = 1000
    # 添加迭代次数
    iter_times = 30
    # 根据迭代次数=需要保存的预测参数文件个数
    # pre_data = [Pre(i) for i in range(epochs)]
    # 记录运行时的所有真实值，VM, 误差
    tru = pd.DataFrame(columns=['tru_rx', 'tru_ry', 'tru_rz', 'tru_tx', 'tru_ty', 'tru_tz'])
    best_pre = pd.DataFrame(columns=['pre_rx', 'pre_ry', 'pre_rz', 'pre_tx', 'pre_ty', 'pre_tz'])
    best_VM_dataframe = pd.DataFrame(columns=['min_VM'])
    # 导入数据分析师
    res_analyse = ResAnalyse(predictor.data_generator, predictor.saving_drr_path + 'statistical_result')
    # 加噪声
    noise = [5, 1, 5, 25, 50, 25]  # 前三维是角度对应的噪声，后三维是平移对应的噪声
    img_name = '预测值对应投影图的预测误差'  # 实际生成的对比DRR图片的标题
    for epoch in range(epochs):
        # 初始化第一次的误差为0
        error_iter = np.array([0, 0, 0, 0, 0, 0])
        # 随机生成一组位姿参数
        # 标准正位：（0±5， 270±1， 90±5）
        rx_noise = np.random.uniform(-noise[0], noise[0])
        ry_noise = np.random.uniform(-noise[1], noise[1])
        rz_noise = np.random.uniform(-noise[2], noise[2])
        # 位移±50
        tx = np.random.uniform(-noise[3], noise[3])
        ty = np.random.uniform(-noise[4], noise[4])
        tz = np.random.uniform(-noise[5], noise[5])
        # 获取加噪后生成的图片
        test_drr = predictor.get_drr(rx_noise, ry_noise, rz_noise, tx, ty, tz)
        # 通过网络获取预测的输出
        pre_r1, pre_t1 = predictor.predict(add_coord(test_drr))
        first_pre = np.array([pre_r1[0], pre_r1[1], pre_r1[2], pre_t1[0], pre_t1[1], pre_t1[2]])
        latest_pre = first_pre
        # 计算第一次的VM
        tru_angle = np.array([rx_noise, 270 + ry_noise, 90 + rz_noise])
        VM_iter = vm_sample(tru_angle, np.array([tx, ty, tz]),
                            pre_r1, pre_t1, 1, pt=predictor.CT_pt, is_label_transform=False)
        # 数据记录
        rx, ry, rz = tru_angle[0], tru_angle[1], tru_angle[2]
        tru.loc[epoch] = [rx, ry, rz, tx, ty, tz]
        # pre_data[epoch].pre_data.loc[0] = first_pre
        # pre_data[epoch].error.loc[0] = error_iter
        # pre_data[epoch].VM.loc[0] = VM_iter
        min_pre_err = 100
        min_VM = 0
        best_iter = 0
        min_pre_drr_iter = None
        min_latest_pre_iter = None
        best_pre_drr = None
        accumulated_err = 0
        for iter_time in tqdm(range(iter_times - 1)):
            # 实际输入网络的投影参数=第一次的预测值-（最新的预测值-第一次的预测值）
            pre_drr_para = first_pre - accumulated_err
            pre_drr_para = para_limit(pre_drr_para)
            # 根据网络的输出投影出对应的DRR
            r_noise = rt2noised_rt(pre_drr_para[:3], mode=predictor.mode)
            rx_noise, ry_noise, rz_noise = r_noise[0], r_noise[1], r_noise[2]
            pre_drr = predictor.get_drr(rx_noise, ry_noise, rz_noise,
                                        pre_drr_para[3], pre_drr_para[4], pre_drr_para[5])
            # 将网络的预测值生成的drr再输入网络，获取第二次的预测值
            pre_r2, pre_t2 = predictor.predict(add_coord(pre_drr))
            latest_pre = np.array([pre_r2[0], pre_r2[1], pre_r2[2], pre_t2[0], pre_t2[1], pre_t2[2]])
            error_iter = latest_pre - first_pre
            accumulated_err += error_iter
            VM_iter = vm_sample(np.array([rx, ry, rz]), np.array([tx, ty, tz]),
                                pre_drr_para[:3], pre_drr_para[3:], 1, pt=predictor.CT_pt, is_label_transform=False)
            # pre_data[epoch].pre_data.loc[iter_time + 1] = latest_pre
            # pre_data[epoch].error.loc[iter_time + 1] = error_iter
            # pre_data[epoch].VM.loc[iter_time + 1] = VM_iter
            if abs(sum(latest_pre - first_pre)) < abs(min_pre_err):
                best_iter = iter_time
                min_pre_err = sum(latest_pre - first_pre)
                min_VM = VM_iter
                min_pre_drr_iter = pre_drr_para
                min_latest_pre_iter = latest_pre
                best_pre_drr = pre_drr
            # print("第{0}轮，第{1}次迭代，预测误差{2}, VM误差{3}".format(epoch, iter_time, sum(latest_pre - first_pre), VM_iter))
        print("第{0}轮，第{1}次迭代最优，最小预测误差{2}, 最小VM误差{3}".format(epoch, best_iter, min_pre_err, min_VM))
        print("真值            ：", np.array([rx, ry, rz, tx, ty, tz]).round(3))
        print("当前投影参数      ：", min_pre_drr_iter.round(3))
        print("对真值的预测      ：", first_pre.round(3))
        print("对当前投影参数的预测：", min_latest_pre_iter.round(3))
        best_VM_dataframe.loc[epoch] = min_VM
        best_pre.loc[epoch] = min_pre_drr_iter
        # 保存每一轮迭代后的最优结果的可视图
        best_visualization()

    # 导入数据
    res_analyse.load_data(tru, best_pre, best_VM_dataframe['min_VM'])
    res_analyse.plot_3D_point(1, predictor.saving_drr_path + "statistical_result/")
    res_analyse.drr_visualization(1, predictor.saving_drr_path + "VM_comparison/", posture=predictor.mode)
    res_analyse.draw_histogram()
    # 数据保存
    res_analyse.to_csv(datas=[tru] +
                             [best_VM_dataframe] +
                             [best_pre],
                       names=['tru.csv'] +
                             ["best_VM.csv"] +
                             ["best_pre.csv"],
                       saving_path=predictor.saving_drr_path + "datasheet/")
    tru_rx = np.array(tru[['tru_rx', 'tru_ry', 'tru_rz']])
    tru_tx = np.array(tru[['tru_tx', 'tru_ty', 'tru_tz']])
    pre_rx = np.array(best_pre[['pre_rx', 'pre_ry', 'pre_rz']])
    pre_tx = np.array(best_pre[['pre_tx', 'pre_ty', 'pre_tz']])
    error_angle = np.abs(tru_rx - pre_rx)
    error_trans = np.abs(tru_tx - pre_tx)
    print("VM均值:{0}".format(mean(best_VM_dataframe['min_VM'])))
    print("角度均值：{0}".format(mean(error_angle, axis=0)))
    print("平移均值：{0}".format(mean(error_trans, axis=0)))
    print("VM方差:{0}".format(var(best_VM_dataframe['min_VM'])))
    print("角度方差：{0}".format(var(error_angle, axis=0)))
    print("平移方差：{0}".format(var(error_trans, axis=0)))
    # res_analyse.to_csv(datas=[i.pre_data for i in pre_data] +
    #                          [i.error for i in pre_data] +
    #                          [i.VM for i in pre_data] +
    #                          [tru] +
    #                          [best_VM_dataframe] +
    #                          [best_pre],
    #                    names=["pre_iter/pre_data{}.csv".format(i) for i in range(epochs)] +
    #                          ["error_iter/error_data{}.csv".format(i) for i in range(epochs)] +
    #                          ["VM_iter/VM_data{}.csv".format(i) for i in range(epochs)] +
    #                          ['tru.csv'] +
    #                          ["best_VM.csv"] +
    #                          ["best_pre.csv"],
    #                    saving_path=predictor.saving_drr_path + "datasheet/")
